package cn.gwm.flink.streaming.source;

import com.alibaba.ververica.connectors.datahub.source.DatahubRecordReader;
import com.alibaba.ververica.connectors.datahub.source.DatahubShardInputSplit;
import com.aliyun.datahub.client.exception.DatahubClientException;
import com.aliyun.datahub.client.model.GetSubscriptionOffsetResult;
import com.aliyun.datahub.client.model.SubscriptionOffset;
import com.aliyun.datahub.clientlibrary.consumer.Consumer;
import com.aliyun.datahub.clientlibrary.exception.ExceptionRetryer;
import com.aliyun.datahub.clientlibrary.models.Offset;
import org.apache.flink.api.common.functions.RuntimeContext;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.core.io.InputSplit;
import org.apache.flink.metrics.Counter;
import org.apache.flink.metrics.SimpleCounter;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * 自定义datahub 读取数据类
 * @ClassName DatahubRecordsReader
 * @Author xzh
 * @Date 2023/03/09 11:10
 **/
public class DatahubRecordsReader extends DatahubRecordReader {

    private static final Logger LOG = LoggerFactory.getLogger(DatahubRecordsReader.class);

    private transient DoubleGauge latencyGauge;
    private transient Counter exceptionCounter;
    private transient Counter readCounter;
    private transient Counter readNullCounter;
    private transient Counter readInvalidCounter;
    //shards 列表
    private List<String> initShardsList;
    public DatahubRecordsReader(String endPoint, String projectName, String topicName, String subId, String accessId, String accessKey, int initShardNum, long startTimeInMs, long stopTimeInMs, List<String> initShardsList) {
        super(endPoint, projectName, topicName, subId, accessId, accessKey, initShardNum, startTimeInMs, stopTimeInMs);
        this.initShardsList = initShardsList;
    }

    public DatahubRecordsReader(String endPoint, String projectName, String topicName, String subId, Configuration properties, int initShardNum, long startTimeInMs, long stopTimeInMs, List<String> initShardsList) {
        super(endPoint, projectName, topicName, subId, properties, initShardNum, startTimeInMs, stopTimeInMs);
        this.initShardsList = initShardsList;
    }
    public void open(InputSplit split, RuntimeContext context) throws IOException {
        this.inputSplit = (DatahubShardInputSplit)split;
        this.initClientProvider();
        if (this.consumerConfig != null) {
            throw new IOException("consumerConfig is not null, not expected!");
        } else {
            this.consumerConfig = this.createConsumerConfig();
            if (this.consumer != null) {
                throw new IOException("consumer is not null, not expected!");
            } else {
                Map<String, Offset> offsetMap = new HashMap();
                Offset offset = new Offset();
                GetSubscriptionOffsetResult offset1 = this.clientProvider.getClient().getSubscriptionOffset(this.projectName, this.topicName, this.subId, this.initShardsList);
                if(offset1.getOffsets().size()>0){
                    SubscriptionOffset subscriptionOffset = offset1.getOffsets().get(this.inputSplit.getShardId());
                    if(subscriptionOffset != null && subscriptionOffset.getTimestamp()>0&&subscriptionOffset.getSequence()>0){
                        offset.setSequence(subscriptionOffset.getSequence());
                        offset.setTimestamp(subscriptionOffset.getTimestamp());
                        LOG.info("Open [{}] using Sequence: {} timestamp:{} ", split.toString(), subscriptionOffset.getSequence(),subscriptionOffset.getTimestamp());
                    }else{
                        offset.setTimestamp(0L);
                        LOG.info("Open [{}] using timestamp:{} ", split.toString(), 0);
                    }
                }else{
                    offset.setTimestamp(0L);
                    LOG.info("Open [{}] using timestamp:{} ", split.toString(), 0);
                }
                offsetMap.put(this.inputSplit.getShardId(), offset);
                this.consumer = this.createConsumer(offsetMap);
                if (context != null) {
                    this.latencyGauge = (DoubleGauge)context.getMetricGroup().addGroup("input").addGroup(this.getName()).gauge("readLatency", new DoubleGauge());
                    this.exceptionCounter = context.getMetricGroup().addGroup("input").addGroup(this.getName()).counter("exceptionCounter", new SimpleCounter());
                    this.readCounter = context.getMetricGroup().addGroup("input").addGroup(this.getName()).counter("readCounter", new SimpleCounter());
                    this.readNullCounter = context.getMetricGroup().addGroup("input").addGroup(this.getName()).counter("nullCounter", new SimpleCounter());
                    this.readInvalidCounter = context.getMetricGroup().addGroup("input").addGroup(this.getName()).counter("invalidCounter", new SimpleCounter());
                }

                this.initPartitionNumsListener();
            }
        }
    }

    private Consumer createConsumer(Map<String, Offset> offsetMap) throws IOException {
        long startTime = System.currentTimeMillis();
        int retryTimes = 3;
        Exception exception = null;
        int i = 0;

        while(i <= retryTimes) {
            try {
                this.consumer = new Consumer(this.projectName, this.topicName, this.subId, offsetMap, this.consumerConfig);
                return this.consumer;
            } catch (DatahubClientException var8) {
                if (ExceptionRetryer.isFatalException(var8)) {
                    throw var8;
                }

                exception = var8;
                LOG.warn("fail to create consumer, retry {} times", i, var8);
                this.sleepAndRetry(startTime);
                ++i;
            }
        }

        throw new DatahubClientException(String.format("create consumer failed after retrying %s times", Integer.valueOf(retryTimes)), exception);
    }

    private String getName() {
        return this.name;
    }
}
